#include "context.h"
#include "common.h"

void Context::GetParameters(params param)
{
    switch (param) {
            // 基础算子参数 logN = 15，Q = 39bit + 17 * 36bit，P = 6 * 41bit
        case params::logN15QP891:
            qVec = { 0x4001b00001, 0xfff9c0001,  0xfff8e0001,  0xfff840001,  0xfff700001,  0xfff640001,
                0x10014c0001, 0xfff4c0001,  0xfff3c0001,  0xfff280001,  0xfff100001,  0xffefe0001,
                0xffee80001,  0x10004a0001, 0x1000500001, 0x1000960001, 0x1000a20001, 0x1000b40001 };
            pVec = { 0x10a19000001, 0x10a3b000001, 0x10a83000001, 0x10b0a000001, 0x10b2e000001, 0x10bc1000001 };
            tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001 };
            break;
            // 自举测试参数 logN = 12, Q = 28 * 54bit, P = 4 * 60bit
        case params::logN12QP1725:
            qVec = {
                0x3ffffffffd6001, 0x20000000244001, 0x200000001ee001, 0x2000000021a001, 0x1ffffffff9c001,
                0x20000000116001, 0x1fffffffe28001, 0x200000000ec001, 0x1fffffffe2a001, 0x200000000e0001,
                0x1fffffffe2e001, 0x1ffffffff38001, 0x1fffffffefa001, 0x1ffffffff72001, 0x1fffffffeba001,
                0x1ffffffff32001, 0x1fffffffe30001, 0x200000000c8001, 0x1ffffffff86001, 0x200000000b2001,
                0x1ffffffffa4001, 0x2000000009e001, 0x2000000002e001, 0x20000000044001, 0x2000000001a001,
                0x20000000094001, 0x1ffffffffb4001, 0x2000000000a001,
            };
            pVec = { 1152921504606830593, 1152921504606748673, 1152921504606683137, 1152921504606601217 };
            tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001, };
            break;
            // SEAL参数 logN = 15, Q = 15 * 55bit, P = 56bit
        case params::logN15QP881:
            qVec = { 
                36028797017456641, 36028797014704129, 36028797014573057, 36028797014376449, 36028797013327873,
                36028797013000193, 36028797012606977, 36028797010444289, 36028797009985537, 36028797005856769,
                36028797005529089, 36028797005135873, 36028797003694081, 36028797003563009, 36028797001138177 
            };
            pVec = { 72057594037338113 };
            tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                        0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001 };
            break;
        case params::logN15QP1725:
            qVec = {
                18014398506729473,9007199247138817,9007199239340033,9007199246417921
                ,9007199240060929,9007199245565953,9007199240847361,9007199242813441
                ,9007199241306113,9007199256248321,9007199247532033,9007199256051713
                ,9007199248777217,9007199267389441,9007199249891329,9007199267258369
                ,9007199250087937,9007199266471937,9007199250481153,9007199264636929
                ,9007199250874369,9007199263260673,9007199251660801,9007199262867457
                ,9007199252119553,9007199261294593,9007199252840449,9007199255658497
            };
            pVec = { 1152921504606584833,1152921504598720513,1152921504597016577,1152921504595968001 };
            tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                     0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001 };
            break;
    }
}

void Context::GetParametersLambda(long logN, int lambda, int dNum)
{
    switch (lambda) {
        case 128: {
            switch(logN) {
                case 16: {
                    if((dNum <= 0) || (dNum > L)) {
                        L = 28;
                        K = 4;
                        T = 8;
                        gamma = 4;
                        qVec = {36028797001138177, 36028796997599233, 36028796996681729, 36028796992749569, 
                                36028796991700993, 36028796990390273, 36028796987637761, 36028796986851329, 
                                36028796984098817, 36028796982263809, 36028796982132737, 36028796980953089, 
                                36028796974661633, 36028796973088769, 36028796972040193, 36028796971253761, 
                                36028796970074113, 36028796967976961, 36028796967190529, 36028796966141953, 
                                36028796965224449, 36028796964569089, 36028796960243713, 36028796952379393, 
                                36028796950675457, 36028796950020097, 36028796949626881, 36028796945170433};
                        pVec = {36028797014376449, 36028797013327873, 36028797010444289, 36028797005856769};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                                0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001};
                        break;
                    }
                    else {
                        L = 32 / (dNum + 1) * dNum;
                        K = 32 / (dNum + 1);
                        T = 12;
                        gamma = 4;
                        GetPrimes(lambda, logN, logp);
                        break;
                    }
                }
                case 15: {
                    if((dNum <= 0) || (dNum > L)) {
                        L = 15;
                        K = 1;
                        T = 5;
                        gamma = 4;
                        qVec = { 
                            36028797017456641, 36028797014704129, 36028797014573057, 36028797014376449, 36028797013327873,
                            36028797013000193, 36028797012606977, 36028797010444289, 36028797009985537, 36028797005856769,
                            36028797005529089, 36028797005135873, 36028797003694081, 36028797003563009, 36028797001138177 
                        };
                        pVec = { 72057594037338113 };
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001, 0xfffffffff2a0001
                            };
                        testparams();
                    }
                    else {
                        L = 16 / (dNum + 1) * dNum;
                        K = 16 / (dNum + 1);
                        T = 8;
                        gamma = 4;
                        GetPrimes(lambda, logN, logp);
                    }
                    break;
                }
                case 14: {
                    if((dNum <= 0) || (dNum > L)) {
                        L = 7;
                        K = 1;
                        T = 4;
                        gamma = 2;
                        qVec = {36028797016178689, 36028797014704129, 36028797014573057, 36028797014376449, 
                                36028797014081537, 36028797013327873, 36028797013098497};
                        pVec = {36028797017456641};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001};
                        break;
                    }
                    else {
                        L = 8 / (dNum + 1) * dNum;
                        K = 8 / (dNum + 1);
                        T = 4;
                        gamma = 2;
                        GetPrimes(lambda, logN, logp);
                    }
                }
                case 13: {
                    if((dNum <= 0) || (dNum > L)) {
                        L = 6;
                        K = 1;
                        T = 4;
                        gamma = 2;
                        qVec = {36028797017571329, 36028797017456641, 36028797017276417, 
                                36028797017014273, 36028797016719361, 36028797016588289};
                        pVec = {36028797018652673};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001};
                        break;
                    }
                    else {
                        L = 4 / (dNum + 1) * dNum;
                        K = 4 / (dNum + 1);
                        T = 4;
                        gamma = 2;
                        GetPrimes(lambda, logN, logp);
                    }
                }
                case 12: {
                    logp = 32;
                    L = 2;
                    K = 1;
                    T = 2;
                    gamma = 1;
                    qVec = {0x200000e001, 0x100006001};// 37 + 32
                    pVec = {0x3ffffea001};// 38
                    tVec = { 0xffffffffffc0001, 0xfffffffff840001};
                    break;
                }
                default: 
                    throw std::domain_error("Invalid value for logN.");}
        }
        break;   
        case 192: {
            switch(logN) {
                case 16: {
                    if((dNum <= 0) || (dNum > L)) {
                        L = 20;
                        K = 2;
                        T = 4;
                        gamma = 2;
                        qVec = {36028797010444289, 36028797005856769, 36028797001138177, 36028796997599233, 36028796996681729, 
                                36028796992749569, 36028796991700993, 36028796990390273, 36028796987637761, 36028796986851329, 
                                36028796984098817, 36028796982263809, 36028796982132737, 36028796980953089, 36028796974661633, 
                                36028796973088769, 36028796972040193, 36028796971253761, 36028796970074113, 36028796967976961};
                        pVec = {36028797014376449, 36028797013327873};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001};
                        break;
                    }
                    else {
                        L = 22 / (dNum + 1) * dNum;
                        K = 22 / (dNum + 1);
                        T = 12;
                        gamma = 4;
                        GetPrimes(lambda, logN, logp);
                        testparams();
                        break;
                    }
                }
                case 15: {
                    if((dNum <= 0) || (dNum > L)) {
                        logp = 50;
                        L = 8;
                        K = 4;
                        T = 8;
                        gamma = 4;
                        qVec = {1125899902124033, 1125899901665281, 1125899899174913, 1125899896160257, 
                                1125899887312897, 1125899886395393, 1125899885740033, 1125899885412353};
                        pVec = {1125899904679937, 1125899903827969, 1125899903500289, 1125899903107073};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
                                0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001};
                    break;
                        break;
                    }
                    else {
                        L = 12 / (dNum + 1) * dNum;
                        K = 12 / (dNum + 1);
                        T = 8;
                        gamma = 4;
                        GetPrimes(lambda, logN, logp);
                        testparams();
                        break;
                    }
                }
                case 14: {
                    if((dNum <= 0) || (dNum > L)) {
                        logp = 50;
                        L = 4;
                        K = 2;
                        T = 4;
                        gamma = 2;
                        qVec = {1125899903827969, 1125899903795201, 1125899903500289, 1125899903107073};
                        pVec = {1125899904679937, 1125899903991809};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001};
                        break;
                    }
                    else {
                        L = 6 - 6 / (dNum + 1);
                        K = 6 / (dNum + 1);
                        T = 4;
                        gamma = 2;
                        GetPrimes(lambda, logN, logp);
                        testparams();
                        break;
                    }
                }
                case 13: {
                    if((dNum <= 0) || (dNum > L)) {
                        logp = 40;
                        L = 3;
                        K = 1;
                        T = 2;
                        gamma = 1;
                        qVec = {1099510890497, 1099510824961, 1099510054913};
                        pVec = {1099511480321};
                        tVec = { 0xffffffffffc0001, 0xfffffffff840001};
                        break;
                    }
                    else {
                        L = 4 - 6 / (dNum + 1);
                        K = 4 / (dNum + 1);
                        T = 4;
                        gamma = 2;
                        GetPrimes(lambda, logN, logp);
                        testparams();
                        break;
                    }
                }
                default: 
                    throw std::domain_error("Invalid value for logN.");
            }
        }     
    }
}

bool isPrime(uint64_t n) {
    if (n <= 1) {
        return false;
    }
    if (n == 2) {
        return true;
    }
    if (n % 2 == 0) {
        return false;
    }
    for (uint64_t i = 3; i <= std::sqrt(n); i += 2) {
        if (n % i == 0) return false;
    }
    
    return true;
}

void Context::GetPrimes(uint64_t lambda, uint64_t logN, uint64_t bit)
{
    N = 1L << logN;
    uint64_t a = (1L << bit) + 1;
    for(int i = 0; i < K; i++) {
        while (true) {
            a -= (2 * N);
            if (isPrime(a)) {
                pVec.push_back(a);
                break;
            }
        } 
    }
    for(int i = 0; i < L; i++) {
        while (true) {
            a -= (2 * N);
            if (isPrime(a)) {
                qVec.push_back(a);
                break;
            }
        }  
    }
    vector<uint64_t> t = { 0xffffffffffc0001, 0xfffffffff840001, 0xfffffffff6a0001, 0xfffffffff5a0001,
             0xfffffffff2a0001, 0xfffffffff240001, 0xffffffffefe0001, 0xffffffffeca0001,
             0xffffffffe9e0001, 0xffffffffe7c0001, 0xffffffffe740001, 0xffffffffe520001, };
    for(int i = 0; i < T; i++) {
        tVec.push_back(t[i]);
    }
             dNum = ceil(L / (K * 1.0));
    testparams();
}

#include <cassert>
#include <NTL/RR.h>
void Context::testparams()
{
    N = 1L << logN;
    dNum = ceil(L / (K * 1.0));
    NTL::ZZ mulP(1);
    NTL::ZZ mulPQ_gamma(1);
    NTL::ZZ mulQ_alpha(1);
    NTL::ZZ mulT(1);
    for(int i = 0; i < K; i++) {
        mulP *= pVec[i];
        mulQ_alpha *= qVec[i];
    }
    for(int i = 0; i < gamma; i++) {
        if(i < K) {
            mulPQ_gamma *= pVec[i];
        }
        else {
            mulPQ_gamma *= qVec[i - K];
        }
    }
    for(int i = 0; i < T; i++) {
        mulT *= tVec[i];
    }
    assert(mulP > mulQ_alpha);
    assert(mulT > (mulQ_alpha * mulPQ_gamma * N * dNum));
}
